package com.lixw.aiassistant.config;

import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.TimeUnit;

/**
 * @ClassName: VectorStoreConfig
 * @Description:
 * @Author: xuweiLi
 * @Create: 2025/10/13 19:41
 **/
@Slf4j
@Configuration
public class VectorStoreConfig {

    @Value("${spring.ai.vectorstore.qdrant.collection-name}")
    private String qdrantCollectionName;

    /**
     * 创建向量存储对象
     * @param qdrantClient
     * @param embeddingModel
     * @return
     */
    @Bean
    public QdrantVectorStore vectorStore(QdrantClient qdrantClient, EmbeddingModel embeddingModel) {
        // 获取 embedding 向量维度
        int vectorSize = embeddingModel.dimensions();
        // 构建 CreateCollection 请求
        Collections.VectorParams vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(vectorSize)
                .build();
        // 发送创建请求（如果 collection 已存在会失败，可捕获异常忽略）
        try {
            qdrantClient.createCollectionAsync(qdrantCollectionName, vectorParams)
                    .get(30, TimeUnit.SECONDS); // 等待完成
            log.info("Qdrant collection '{}' created with Cosine distance.", qdrantCollectionName);
        } catch (Exception e) {
            // 如果只是“已存在”，可以忽略
            log.warn(" Collection may already exist!");
        }
        return QdrantVectorStore.builder(qdrantClient, embeddingModel)
                .collectionName(qdrantCollectionName)
                .initializeSchema(false)
                .batchingStrategy(new TokenCountBatchingStrategy())
                .build();
    }

    @Bean
    public VectorStoreDocumentRetriever documentRetriever(QdrantVectorStore vectorStore) {
        return VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .topK(5)
                .similarityThreshold(0.2)
                .build();
    }
}